import pytorch_lightning as L
import torch

from lib.metrics import update_miou_metrics
from .loss import structure_loss


class BaseModule(L.LightningModule):

    def __init__(self, num_classes):
        super(BaseModule, self).__init__()
        self.num_classes = num_classes

    def training_step(self, batch, batch_idx, train_metrics, normalization):
        data, label, _ = batch
        bg_mask = 1-label
        lateral_map_2_fg,lateral_map_3_fg,lateral_map_4_fg,lateral_map_5_fg,lateral_map_2_bg,lateral_map_3_bg,lateral_map_4_bg,lateral_map_5_bg = self(data)
        loss5 = structure_loss(lateral_map_2_fg, lateral_map_2_bg, label, bg_mask, self.num_classes)
        loss4 = structure_loss(lateral_map_3_fg, lateral_map_3_bg, label, bg_mask, self.num_classes)
        loss3 = structure_loss(lateral_map_4_fg, lateral_map_4_bg, label, bg_mask, self.num_classes)
        loss2 = structure_loss(lateral_map_5_fg, lateral_map_5_bg, label, bg_mask, self.num_classes)
        loss = loss2 + loss3 + loss4 + loss5
        predict = normalization(lateral_map_5_fg)
        update_miou_metrics(train_metrics, predict, label)
        return loss

    def validation_step(self, batch, batch_idx, valid_metrics, normalization):
        data, label, _ = batch
        bg_mask = 1-label
        lateral_map_2_fg,lateral_map_3_fg,lateral_map_4_fg,lateral_map_5_fg,lateral_map_2_bg,lateral_map_3_bg,lateral_map_4_bg,lateral_map_5_bg = self(data)
        predict = normalization(lateral_map_5_fg)
        update_miou_metrics(valid_metrics, predict, label)
        return structure_loss(lateral_map_5_fg, lateral_map_5_bg, label, bg_mask, self.num_classes)
    
    def test_step(self, batch, batch_idx):
        img, msk, _ = batch
        lateral_map_2_fg,lateral_map_3_fg,lateral_map_4_fg,lateral_map_5_fg,lateral_map_2_bg,lateral_map_3_bg,lateral_map_4_bg,lateral_map_5_bg = self(img)
        return torch.squeeze(lateral_map_5_fg, dim=1)
